#!/usr/bin/env python3

import rospy
import math
import tf2_ros
import tf2_geometry_msgs
from nav_msgs.msg import OccupancyGrid
from geometry_msgs.msg import PoseStamped
from tf.transformations import euler_from_quaternion


class ClipLocalMap:
    def __init__(self):
        rospy.init_node("clip_local_map_node")

        self.global_map = None
        self.listener = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.listener)

        rospy.Subscriber("/map", OccupancyGrid, self.map_callback)
        self.local_map_pub = rospy.Publisher("/local_map_clip", OccupancyGrid, queue_size=10)

        self.timer = rospy.Timer(rospy.Duration(0.1), self.timer_callback) 

    def map_callback(self, msg):
        self.global_map = msg

    def timer_callback(self, event):
        if self.global_map is None:
            return

        try:
            trans = self.listener.lookup_transform("map", "body", rospy.Time(0), rospy.Duration(0.1))
            x = trans.transform.translation.x
            y = trans.transform.translation.y
            q = trans.transform.rotation
            _, _, yaw = euler_from_quaternion([q.x, q.y, q.z, q.w])
        except (tf2_ros.LookupException, tf2_ros.ExtrapolationException):
            rospy.logwarn("TF transform unavailable")
            return

        self.publish_local_map(x, y, yaw)

    def publish_local_map(self, x, y, yaw):
        resolution = self.global_map.info.resolution
        map_width = self.global_map.info.width
        map_height = self.global_map.info.height
        origin_x = self.global_map.info.origin.position.x
        origin_y = self.global_map.info.origin.position.y

        width = int(12.0 / resolution)
        height = int(12.0 / resolution)

        local_map = OccupancyGrid()
        local_map.header.stamp = rospy.Time.now()
        local_map.header.frame_id = "body"
        local_map.info.resolution = resolution
        local_map.info.width = width
        local_map.info.height = height
        local_map.info.origin.position.x = -6.0
        local_map.info.origin.position.y = -6.0
        local_map.info.origin.orientation.w = 1.0

        data = []
        for row in range(height):
            for col in range(width):
                dx = (col + 0.5) * resolution - 6.0
                dy = (row + 0.5) * resolution - 6.0

                x_map = x + dx * math.cos(yaw) - dy * math.sin(yaw)
                y_map = y + dx * math.sin(yaw) + dy * math.cos(yaw)

                cell_x = int((x_map - origin_x) / resolution)
                cell_y = int((y_map - origin_y) / resolution)

                if 0 <= cell_x < map_width and 0 <= cell_y < map_height:
                    idx = cell_y * map_width + cell_x
                    value = self.global_map.data[idx]
                else:
                    value = -1
                data.append(value)

        local_map.data = data
        self.local_map_pub.publish(local_map)


if __name__ == "__main__":
    try:
        ClipLocalMap()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
